Add `sym_sizes` with dynamic shape support #3909

miladm · 2022-08-19T20:52:26Z

Add sym_sizes with dynamic shape support

miladm · 2022-08-19T20:54:27Z

This PR shows that this implementation must be correct. However, there are other issues we want to resolves in the mentioned PR that may require extra time. This PR proposes the landing of the sym_sizes independently given all we know about its implementation methodology.

Krovatkin · 2022-09-13T21:49:03Z

@JackCaoG I remember you had some doubts about this PR?

Krovatkin · 2022-09-13T23:52:15Z

never mind, I'm seeing real failures now, need to investigate @JackCaoG

JackCaoG · 2022-09-13T23:53:53Z

seems OK from high level, I am not sure what does this pr enabled us through. is it .sym_size for XLATensorImpl?

vanbasten23 · 2022-09-16T18:21:23Z

torch_xla/csrc/tensor_impl.cpp

  return sizes_default();
 }

 c10::SymIntArrayRef XLATensorImpl::sym_sizes_custom() const {


Wondering what the "custom" in sym_sizes_custom() implies

custom means that a given TensorImpl can implement sym_sizes in whichever way they see fit. Note, sym_sizes_custom is a virtual method whereas sym_sizes() isn't. sym_sizes() checks if a given tensor impl implements CustomSizes policy and if so calls a virtual sym_sizes. This is kinda convoluted and was done to make sure we don't pay a virtual call penalty on the fast path: cpu and cuda tensors don't need the custom_sym_sizes implementation.

vanbasten23 · 2022-09-16T18:43:25Z

torch_xla/csrc/tensor_impl.cpp

+        sym_sizes_.push_back(sn);
+        /*TODO(miladm): verify numel_ calculation after adding a dynamic op
+         */
+        numel_ *= dynamic_cast<SizeNode*>(dim_node.get())->getStaticValue();


I've read the definition of numel_ in the base class but it's not clear to me what it means. Do you know the meaning of this member?

number of elements

Krovatkin · 2022-09-19T22:01:35Z

seems OK from high level, I am not sure what does this pr enabled us through. is it .sym_size for XLATensorImpl?

This PR finally opens the Pandora box, hehe. For ops like nonzero size() or shape in python will now return real sym ints and for any code in C++ that calls sym_sizes().

test/cpp/test_aten_xla_tensor.cpp

JackCaoG · 2022-09-19T23:52:59Z

test/cpp/test_aten_xla_tensor.cpp

+    ASSERT_EQ(a.sym_sizes().at(0).is_symbolic(), false);
+  });
+}
+


should we add a dynamic test? You can call non_zero but you would need to make sure https://github.com/pytorch/xla/blob/master/test/cpp/run_tests.sh#L11 is set so non_zero won't fallback to cpu.

JackCaoG · 2022-09-20T17:45:05Z

torch_xla/csrc/tensor_impl.h

+#include <torch/csrc/lazy/backend/backend_interface.h>
+#include <torch/csrc/lazy/core/config.h>
+#include <torch/csrc/lazy/core/ir.h>
+#include <torch/csrc/lazy/core/trie.h>


do we still needs these includes?

miladm · 2022-09-27T19:30:03Z

torch_xla/csrc/tensor_impl.cpp

    sizes_and_strides_.set_sizes(updated_sizes);
    auto updated_strides = torch::lazy::ComputeArrayStrides(
        torch::lazy::ToVector<int64_t>(shape.get().dimensions()));
    for (int i = 0; i < updated_strides.size(); i++) {


Shouldn't we move this for loop inside SetupSymSizeProperties()? This makes SetupSizeProperties and SetupSymSizeProperties implementations consistent. Is there a reason we want them called outside SetupSymSizeProperties?

@Krovatkin

if I understand your point, you are suggesting to duplicate the strides logic inside SetupSymSizeProperties as well, so it updates both sym_sizes and sym_strides the same way SetupSizeProperties updates sizes and strides?

If this is your concern, I believe it's fine for now, since sym_strides is still returning static strides now.
We should figure out if we need support for dynamic strides as well or whether it's okay for sym_strides to continue returning static strides.

@ezyang: based on our offline conversation about sym_strides under functionalization, it seems we assume no need for dynamic strides; correct?

miladm · 2022-09-27T19:30:15Z

torch_xla/csrc/tensor_impl.cpp

      sizes_and_strides_.stride_at_unchecked(i) = updated_strides[i];
    }
+    SetupSymSizeProperties();
    generation_ = generation;


miladm · 2022-09-27T19:35:42Z

test/cpp/test_aten_xla_tensor.cpp

+    xla_b = torch::nonzero(xla_b);
+    auto s0 = xla_b.sym_sizes().at(0);
+    ASSERT_EQ(s0.is_symbolic(), true);
+    auto sininode =


nit: is the variable name meant to be sininode?

@Krovatkin

miladm requested review from JackCaoG and Krovatkin August 19, 2022 20:52

miladm self-assigned this Aug 19, 2022

miladm added the dynamism Dynamic Shape Features label Aug 19, 2022

miladm added this to the Dynamic Shape milestone Aug 19, 2022

Krovatkin force-pushed the sym_sizes_dynamic branch 2 times, most recently from ceb2954 to 149d61e Compare September 13, 2022 21:46

vanbasten23 reviewed Sep 16, 2022

View reviewed changes

Krovatkin force-pushed the sym_sizes_dynamic branch from 149d61e to 06e9ecf Compare September 19, 2022 21:56

JackCaoG reviewed Sep 19, 2022

View reviewed changes

test/cpp/test_aten_xla_tensor.cpp Show resolved Hide resolved

JackCaoG reviewed Sep 19, 2022

View reviewed changes

JackCaoG reviewed Sep 20, 2022

View reviewed changes

Krovatkin requested review from JackCaoG and vanbasten23 September 25, 2022 19:29

Krovatkin force-pushed the sym_sizes_dynamic branch 2 times, most recently from 3954208 to 89062bd Compare September 26, 2022 18:02

miladm and others added 5 commits September 27, 2022 16:54

add sym_sizes with dynamic shape support

8a0d84f

test fixes

5d9451c

remove an unnecessary cast

051ecf9

fix test failures

1deb416

addderss Jack's feedback

e4befa3

Krovatkin force-pushed the sym_sizes_dynamic branch from 89062bd to e4befa3 Compare September 27, 2022 16:54

miladm commented Sep 27, 2022

View reviewed changes

JackCaoG approved these changes Sep 28, 2022

View reviewed changes

JackCaoG merged commit 0efa9e4 into master Sep 28, 2022

Add sym_sizes with dynamic shape support #3909

Add sym_sizes with dynamic shape support #3909

Uh oh!

Conversation

miladm commented Aug 19, 2022

Uh oh!

miladm commented Aug 19, 2022

Uh oh!

Krovatkin commented Sep 13, 2022

Uh oh!

Krovatkin commented Sep 13, 2022

Uh oh!

JackCaoG commented Sep 13, 2022

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Krovatkin commented Sep 19, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

miladm Sep 27, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Krovatkin Sep 27, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

miladm Sep 27, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Add `sym_sizes` with dynamic shape support #3909

Add `sym_sizes` with dynamic shape support #3909

Krovatkin commented Sep 19, 2022 •

edited

Loading

miladm Sep 27, 2022 •

edited

Loading

Krovatkin Sep 27, 2022 •

edited

Loading

miladm Sep 27, 2022 •

edited

Loading